internal/handler/document_test.go (2)

480-539: ⚡ Quick win

Add tests for Content-Disposition headers and parameter validation.

The test coverage for GetDocumentArtifact is missing:

1. Validation of the Content-Disposition header. According to the PR objectives, there's attachment handling logic (shouldForceArtifactAttachment). The success test should verify that when ForceAttachment is false, the header is set to inline; filename="safe.txt", and there should be a test case where ForceAttachment is true to verify the attachment disposition.
2. Edge case tests for parameter validation (e.g., empty or missing filename parameter).

📋 Suggested test additions

Add a test to verify the Content-Disposition header in the success case:

 func TestGetDocumentArtifact_Success(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	h := &DocumentHandler{
 		documentService: &fakeDocumentService{},
 	}
 	c, w := setupGinContextWithUser("GET", "/api/v1/documents/artifact/test.txt", "")
 	c.Params = gin.Params{{Key: "filename", Value: "test.txt"}}
 
 	h.GetDocumentArtifact(c)
 
 	if w.Code != http.StatusOK {
 		t.Fatalf("expected 200, got %d", w.Code)
 	}
 	if w.Header().Get("Content-Type") != "text/plain" {
 		t.Fatalf("unexpected content type: %s", w.Header().Get("Content-Type"))
 	}
+	if contentDisp := w.Header().Get("Content-Disposition"); contentDisp != "inline; filename=\"safe.txt\"" {
+		t.Fatalf("unexpected Content-Disposition: %s", contentDisp)
+	}
 	if w.Body.String() != "artifact content" {
 		t.Fatalf("unexpected body: %s", w.Body.String())
 	}
 }

Add a test for the attachment disposition case by extending the fake service to return a response with ForceAttachment: true for a specific filename:

func TestGetDocumentArtifact_Attachment(t *testing.T) {
	gin.SetMode(gin.TestMode)
	// Create a custom fake that returns ForceAttachment=true
	fake := &fakeDocumentService{}
	// Override the method for this specific test, or extend fake to handle "attachment.txt"
	h := &DocumentHandler{
		documentService: fake,
	}
	c, w := setupGinContextWithUser("GET", "/api/v1/documents/artifact/attachment.txt", "")
	c.Params = gin.Params{{Key: "filename", Value: "attachment.txt"}}

	h.GetDocumentArtifact(c)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	if contentDisp := w.Header().Get("Content-Disposition"); !strings.HasPrefix(contentDisp, "attachment;") {
		t.Fatalf("expected attachment disposition, got: %s", contentDisp)
	}
}

Add a test for empty filename:

func TestGetDocumentArtifact_EmptyFilename(t *testing.T) {
	gin.SetMode(gin.TestMode)
	h := &DocumentHandler{
		documentService: &fakeDocumentService{},
	}
	c, w := setupGinContextWithUser("GET", "/api/v1/documents/artifact/", "")
	c.Params = gin.Params{{Key: "filename", Value: ""}}

	h.GetDocumentArtifact(c)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	var resp map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &resp)
	if resp["code"] == float64(common.CodeSuccess) {
		t.Fatal("expected error for empty filename")
	}
}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/handler/document_test.go` around lines 480 - 539, Update tests for
GetDocumentArtifact on DocumentHandler to assert Content-Disposition behavior
and validate filename params: extend or modify fakeDocumentService to return
ForceAttachment=true for a specific filename (e.g., "attachment.txt") and add a
TestGetDocumentArtifact_Attachment that calls
DocumentHandler.GetDocumentArtifact and asserts
w.Header().Get("Content-Disposition") starts with "attachment;"; also enhance
the success test (TestGetDocumentArtifact_Success) to assert Content-Disposition
equals `inline; filename="safe.txt"` when ForceAttachment is false (use the fake
to return that behavior), and add a TestGetDocumentArtifact_EmptyFilename that
calls GetDocumentArtifact with an empty gin Param filename and asserts the
response is an error (non-success code) instead of CodeSuccess; reference
DocumentHandler, GetDocumentArtifact, and fakeDocumentService to locate the code
to change.

---

541-621: ⚡ Quick win

Add tests for Content-Disposition headers and dataset mismatch.

The test coverage is missing:

1. Validation of Content-Disposition headers for both GetDocumentPreview and DownloadDocument. The PR objectives mention these endpoints set appropriate headers.
2. For DownloadDocument, a critical test case for dataset mismatch: when a document exists but belongs to a different dataset than the one specified in the URL. According to the PR objectives, the service "enforce[s] strict dataset alignment and access checks."
3. Parameter validation edge cases (empty docID, empty datasetID).

📋 Suggested test additions

Add Content-Disposition validation to the success tests:

 func TestGetDocumentPreview_Success(t *testing.T) {
 	// ... existing setup ...
 	h.GetDocumentPreview(c)
 
 	if w.Code != http.StatusOK {
 		t.Fatalf("expected 200, got %d", w.Code)
 	}
 	if w.Header().Get("Content-Type") != "text/plain" {
 		t.Fatalf("unexpected content type: %s", w.Header().Get("Content-Type"))
 	}
+	if contentDisp := w.Header().Get("Content-Disposition"); contentDisp != "inline; filename=\"preview.txt\"" {
+		t.Fatalf("unexpected Content-Disposition: %s", contentDisp)
+	}
 	if w.Body.String() != "preview content" {
 		t.Fatalf("unexpected body: %s", w.Body.String())
 	}
 }

Add a dataset mismatch test for DownloadDocument:

func TestDownloadDocument_DatasetMismatch(t *testing.T) {
	gin.SetMode(gin.TestMode)
	// Extend fakeDocumentService to return an error for dataset mismatch
	// or use a specific combination that triggers this error
	fake := &fakeDocumentService{}
	h := &DocumentHandler{
		documentService: fake,
	}
	c, w := setupGinContextWithUser("GET", "/api/v1/datasets/wrong-dataset/documents/doc-1", "")
	c.Params = gin.Params{{Key: "dataset_id", Value: "wrong-dataset"}, {Key: "document_id", Value: "doc-1"}}

	h.DownloadDocument(c)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	var resp map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &resp)
	// Verify appropriate error code for dataset mismatch
	if resp["code"] == float64(common.CodeSuccess) {
		t.Fatal("expected error for dataset mismatch")
	}
}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/handler/document_test.go` around lines 541 - 621, Add assertions for
Content-Disposition headers in TestGetDocumentPreview_Success (verify the header
includes an appropriate filename and disposition like "inline; filename=...")
and in TestDownloadDocument_Success (verify "attachment; filename=..."), extend
fakeDocumentService or add a specific behavior in the fake used by
DownloadDocument tests so that requesting doc-1 with a different dataset_id
returns a dataset-mismatch error and write a new
TestDownloadDocument_DatasetMismatch that calls DownloadDocument with dataset_id
"wrong-dataset" and asserts the response error code is not common.CodeSuccess,
and finally add lightweight parameter validation tests calling
GetDocumentPreview and DownloadDocument with empty docID and/or datasetID (use
setupGinContextWithUser and setting c.Params to empty "") expecting the handler
to return the appropriate error code; reference GetDocumentPreview,
DownloadDocument, fakeDocumentService, setupGinContextWithUser in your changes.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes